
# ========================================================== # 
# v0.0.2 20241104
# add more options for the wavelength 
# add a function that get the optical parameters from the json file 
# add extra json file (mcx_parameter.json)

# v0.0.3 20241111
# modify save options in stage2 run , now the save nii.gz file works well 


# ========================================================== # 


import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
import pmcx
from scipy.ndimage import zoom
import copy
import os
import argparse

def show_src_vol(src_pos,src_dir,vol):
    src_pos=src_pos.astype(int)
    # Extract the x, y, z planes at the source position
    x_pos, y_pos, z_pos = src_pos[0], src_pos[1], src_pos[2]
    x_plane = vol[x_pos, :, :]
    y_plane = vol[:, y_pos, :]
    z_plane = vol[:, :, z_pos]
    # Plotting in 3 axes
    plt.figure(figsize=(15, 5))
    # Plot for x-y plane
    plt.subplot(1, 3, 1)
    plt.imshow(x_plane.T, cmap='gray', origin='lower', aspect='equal')
    plt.colorbar(label='Tissue Label')
    plt.plot(y_pos, z_pos, 'ro', markersize=10, label='Source')
    arrow_scale = 20
    arrow_y = y_pos + arrow_scale * src_dir[1]
    arrow_z = z_pos + arrow_scale * src_dir[2]
    plt.arrow(y_pos, z_pos, arrow_y - y_pos, arrow_z - z_pos, color='red', head_width=5, head_length=5, label='Direction')
    plt.xlabel('y [voxels]')
    plt.ylabel('z [voxels]')
    plt.title('X-Y Plane with Source Position and Direction')
    plt.legend(loc='upper right')
    # Plot for y-z plane
    plt.subplot(1, 3, 2)
    plt.imshow(y_plane, cmap='gray', origin='lower', aspect='equal')
    plt.colorbar(label='Tissue Label')
    plt.plot(z_pos, x_pos, 'ro', markersize=10, label='Source')
    arrow_y = z_pos + arrow_scale * src_dir[2]
    arrow_x = x_pos + arrow_scale * src_dir[0]
    plt.arrow(z_pos, x_pos, arrow_y - z_pos, arrow_x - x_pos, color='red', head_width=5, head_length=5, label='Direction')
    plt.xlabel('z [voxels]')
    plt.ylabel('x [voxels]')
    plt.title('Y-Z Plane with Source Position and Direction')
    plt.legend(loc='upper right')
    # Plot for x-z plane
    plt.subplot(1, 3, 3)
    plt.imshow(z_plane.T, cmap='gray', origin='lower', aspect='equal')
    plt.colorbar(label='Tissue Label')
    plt.plot(x_pos, y_pos, 'ro', markersize=10, label='Source')
    arrow_x = x_pos + arrow_scale * src_dir[0]
    arrow_y = y_pos + arrow_scale * src_dir[1]
    plt.arrow(x_pos, y_pos, arrow_x - x_pos, arrow_y - y_pos, color='red', head_width=5, head_length=5, label='Direction')
    plt.xlabel('x [voxels]')
    plt.ylabel('y [voxels]')
    plt.title('X-Z Plane with Source Position and Direction')
    plt.legend(loc='upper right')
    
    # Adjust layout and show the plot
    plt.tight_layout()
    plt.show()
def view_result(res,vol):
    CWfluence = np.sum(res['flux'], axis=3)
    #CWfluence*=mask
    old_vol=vol
    new_vol=old_vol#*mask
    
    plt.subplot(2,3,1)
    best_slice_x=np.sum(CWfluence,axis=(1,2)).argmax()
    plt.imshow(np.flipud(np.log10(new_vol[best_slice_x, :, :])), cmap="gray")
    plt.imshow(np.flipud(np.log10(CWfluence[best_slice_x, :, :])), cmap="jet",alpha=0.5)
    plt.title(f"x:{best_slice_x}")
    plt.subplot(2,3,2)
    best_slice_y=np.sum(CWfluence,axis=(0,2)).argmax()
    plt.imshow(np.flipud(np.log10(new_vol[:, best_slice_y, :])), cmap="gray")
    plt.imshow(np.flipud(np.log10(CWfluence[:, best_slice_y, :])), cmap="jet",alpha=0.5)
    plt.title(f"y:{best_slice_y}")
    plt.subplot(2,3,3)
    best_slice_z=np.sum(CWfluence,axis=(0,1)).argmax()
    plt.imshow(np.flipud(np.log10(new_vol[:, :, best_slice_z])), cmap="gray")
    plt.imshow(np.flipud(np.log10(CWfluence[:, :, best_slice_z])), cmap="jet",alpha=0.5)
    plt.title(f"z:{best_slice_z}")

    plt.subplot(2,3,4)
    plt.imshow(np.flipud(old_vol[best_slice_x, :, :]), cmap="gray")
    plt.subplot(2,3,5)
    plt.imshow(np.flipud(old_vol[:, best_slice_y, :]), cmap="gray")
    plt.subplot(2,3,6)
    plt.imshow(np.flipud(old_vol[:, :, best_slice_z]), cmap="gray")
    plt.show()
def view_result_mask(res,vol,mask):
    CWfluence = np.sum(res['flux'], axis=3)
    CWfluence*=mask
    old_vol=vol
    new_vol=old_vol*mask
    
    plt.subplot(2,3,1)
    best_slice_x=np.sum(CWfluence,axis=(1,2)).argmax()
    plt.imshow(np.flipud(np.log10(new_vol[best_slice_x, :, :])), cmap="gray")
    plt.imshow(np.flipud(np.log10(CWfluence[best_slice_x, :, :])), cmap="jet",alpha=0.5)
    plt.title(f"x:{best_slice_x}")
    plt.subplot(2,3,2)
    best_slice_y=np.sum(CWfluence,axis=(0,2)).argmax()
    plt.imshow(np.flipud(np.log10(new_vol[:, best_slice_y, :])), cmap="gray")
    plt.imshow(np.flipud(np.log10(CWfluence[:, best_slice_y, :])), cmap="jet",alpha=0.5)
    plt.title(f"y:{best_slice_y}")
    plt.subplot(2,3,3)
    best_slice_z=np.sum(CWfluence,axis=(0,1)).argmax()
    plt.imshow(np.flipud(np.log10(new_vol[:, :, best_slice_z])), cmap="gray")
    plt.imshow(np.flipud(np.log10(CWfluence[:, :, best_slice_z])), cmap="jet",alpha=0.5)
    plt.title(f"z:{best_slice_z}")

    plt.subplot(2,3,4)
    plt.imshow(np.flipud(old_vol[best_slice_x, :, :]), cmap="gray")
    plt.subplot(2,3,5)
    plt.imshow(np.flipud(old_vol[:, best_slice_y, :]), cmap="gray")
    plt.subplot(2,3,6)
    plt.imshow(np.flipud(old_vol[:, :, best_slice_z]), cmap="gray")
    plt.show()
def from_csv_get_subject_coord(subject_csv,region_name):
    import pandas as pd
    df=pd.read_csv(subject_csv,header=None)
    x,y,z=df[df[4]==region_name].values[0][1:4]
    return np.array([x,y,z])
def run_fn(inputs_dict):
    path,subject_csv,seg_path,region_name,src_dir_mode=inputs_dict["path"],inputs_dict["subject_csv"],inputs_dict["seg_path"],inputs_dict["region_name"],inputs_dict["src_dir_mode"]
    Fpz_subject=from_csv_get_subject_coord(subject_csv,region_name)
    print(inputs_dict,Fpz_subject)
    img = nib.load(path)
    vol=img.get_fdata()[...,0].astype('uint8')

    affine_matrix = img.affine
    print(affine_matrix.shape)
    image_coord = np.linalg.inv(affine_matrix) @ np.append(Fpz_subject, 1)
    src_pos = image_coord[:3]
    print("numpy space",src_pos,vol.shape)
    #src_pos = src_pos * pixdim[:3]
    pixdim = img.header.get_zooms()
    zoom_factor=np.array(pixdim[:3]) / np.array([1.0, 1.0, 1.0])
    vol = zoom(vol, zoom_factor, order=0)
    src_pos = src_pos*zoom_factor
    print("real spcae",src_pos,vol.shape,set(vol.flatten()))
    
    src_dir = None
    target_position = None
    if src_dir_mode=="fixed": # min distance(src_pos) # target(src_pos,tar_pos)
        src_dir=np.array([1.,0.,0.])
        print("using fixed direction, ",src_dir)
    elif src_dir_mode=="default":
        print("using default mode, ignoring seg_path")
        cortex_coords = np.argwhere(np.logical_or(vol==1,vol==2))
        distances = np.linalg.norm(cortex_coords - src_pos, axis=1)
        num_top_points = int(len(distances) * 0.05)
        top_indices = np.argsort(distances)[:num_top_points]
        top_coords = cortex_coords[top_indices]
        target_position = np.mean(top_coords, axis=0)
        direction = target_position - src_pos
        normed_direction = direction / np.linalg.norm(direction)
        src_dir=normed_direction
        print("tar_pos real space",target_position)
        # temp_vol=copy.deepcopy(vol).astype(np.float64)
        # temp_vol[top_indices]=np.nan
        show_src_vol(target_position,-src_dir,vol)
    elif src_dir_mode=="target":
        seg_img=nib.load(seg_path)
        seg_vol=seg_img.get_fdata()
        seg_pixdim = seg_img.header.get_zooms()
        seg_zoom_factor=np.array(seg_pixdim[:3]) / np.array([1.0, 1.0, 1.0])
        seg_vol = zoom(seg_vol, seg_zoom_factor, order=0)
        cortex_coords=np.argwhere(seg_vol==17)# Right # 53.5Left
        top_coords = cortex_coords#[top_indices]
        target_position = np.mean(top_coords, axis=0)
        direction = target_position - src_pos
        normed_direction = direction / np.linalg.norm(direction)
        src_dir=normed_direction
        print("tar_pos real space",target_position)
        print("tar_pos numpy space", target_position/seg_zoom_factor)
        # temp_vol=copy.deepcopy(vol).astype(np.float64)
        # temp_vol[top_coords]=np.nan
        show_src_vol(target_position,-src_dir,vol)
    
    show_src_vol(src_pos,src_dir,vol)
    #show_src_vol(target_position,src_dir,vol)
    
    print(src_pos,src_dir,vol.shape)
    return {"src_pos":src_pos,"src_dir":src_dir,"vol":vol,"target_position":target_position}

def run_fn_stage2(input_dict_v2):
    vol = input_dict_v2["vol"]
    src_dir=input_dict_v2["src_dir"]
    src_pos=input_dict_v2["src_pos"]
    stage_2_mode=input_dict_v2["stage_2_mode"]
    custom_src=input_dict_v2["custom_src"]
    save_path=input_dict_v2["save_path"]
    p=input_dict_v2["p"]
    t=input_dict_v2["t"]
    # custom_src={
    #     "srctype":"disk",
    #     "srcparam1":[12,0,0,0],
    #     "srcparam2":[0,0,0,0],
    # }

    import pmcx
    import nibabel
    import numpy as np
    import nibabel as nib
    import matplotlib.pyplot as plt
    import scipy.io as sio
    import pmcx
    
    ambient_air = np.zeros_like(vol)
    individual_atlas=np.zeros_like(vol)
    # map scalp,csf,gm,wm,air to 1,2,3,4 and rest to 0
    # prop=[
    #   [0, 0, 1, 1], # air
    #   [0.092,38,0.89,1.37], # WM
    #   [0.028,7.3,0.89,1.37], # GM
    #   [0.0026,0.091,0.89,1.37],# CSF
    #   [0,0,1,1], # #4 things
    #   [0.0168,17.82,0.89,1.37],# Scalp
    #   [0,0,1,1],# Eye_balss assuming as air
    #   [0.011,17.82,0.89,1.37],# Compact_bone
    #   [0.011,17.82,0.89,1.37],# Spongy_bone
    #   [0,0,1,1],# Blood assuming as air
    #   [0,0,1,1],# Muscle assuming as air
    # ]
    individual_atlas[vol >0] = 1 # as scalp
    individual_atlas[vol == 5] = 1 # scalp
    individual_atlas[vol == 7] = 2 # skull
    individual_atlas[vol == 8] = 2 # skull
    individual_atlas[vol == 3] = 3 # csf
    individual_atlas[vol == 2] = 4 # gm
    individual_atlas[vol == 1] = 5 # wm 
    individual_atlas[vol >5] = 1 # air 
    
    # Light parameters
    #length = 1064
    wavelength = input_dict_v2["wavelength"]  # 从输入字典获取波长
    length = int(wavelength.replace('nm', ''))
    d = 1
    #p = 250 # 辐照度 mW/
    #t = 8 # mins
    timwin = 1e-9
    # ========================================== #
    # | 波长 | 参数 | 头皮 | 颅骨 | 脑脊液 | 灰质 | 白质 |
    # |------|------|------|------|--------|------|------|
    # | 670nm | μs | 22.65 | 10.82 | 0.091 | 8.4 | 40.1 |
    # | | μa | 0.05625 | 0.0208 | 0.0004 | 0.02 | 0.07 |
    # | | n | 1.37 | 1.37 | 1.37 | 1.37 | 1.37 |
    # | | g | 0.89 | 0.89 | 0.89 | 0.9 | 0.85 |
    # | 810nm | μs | 17.42 | 17.42 | 0.091 | 7.3 | 38 |
    # | | μa | 0.0168 | 0.011 | 0.0026 | 0.028 | 0.092 |
    # | | n | 1.37 | 1.37 | 1.37 | 1.37 | 1.37 |
    # | | g | 0.89 | 0.89 | 0.89 | 0.91 | 0.87 |
    # | 1064nm | μs | 14.63 | 14.63 | 0.091 | 5.9 | 30 |
    # | | μa | 0.019 | 0.019 | 0.0144 | 0.053 | 0.105 |
    # | | n | 1.37 | 1.37 | 1.37 | 1.37 | 1.37 |
    # | | g | 0.89 | 0.89 | 0.89 | 0.91 | 0.88 |
    # ========================================== #
    import json
    import numpy as np

    def get_optical_parameters(wavelength):
        """
        从JSON文件读取指定波长的光学参数
        
        参数:
        wavelength (str): 波长，必须是 '670nm', '810nm' 或 '1064nm'
        
        返回:
        numpy.ndarray: 包含各组织光学参数的数组，顺序为 [μa, μs, g, n]
        """
        # 读取JSON文件
        with open('mcx_parameter.json', 'r') as f:
            parameters = json.load(f)
        
        # 获取指定波长的参数
        tissue_params = parameters['tissue_parameters'][wavelength]
        
        # 创建参数数组，按照原始顺序排列
        optical_params = np.array([
            [0, 0, 1, 1],  # air
            [tissue_params['scalp']['μa'], tissue_params['scalp']['μs'], 
            tissue_params['scalp']['g'], tissue_params['scalp']['n']],  # scalp
            [tissue_params['skull']['μa'], tissue_params['skull']['μs'], 
            tissue_params['skull']['g'], tissue_params['skull']['n']],  # skull
            [tissue_params['CSF']['μa'], tissue_params['CSF']['μs'], 
            tissue_params['CSF']['g'], tissue_params['CSF']['n']],  # csf
            [tissue_params['grey_matter']['μa'], tissue_params['grey_matter']['μs'], 
            tissue_params['grey_matter']['g'], tissue_params['grey_matter']['n']],  # gm
            [tissue_params['white_matter']['μa'], tissue_params['white_matter']['μs'], 
            tissue_params['white_matter']['g'], tissue_params['white_matter']['n']],  # wm # air cavities
        ])
        
        return optical_params

    # 使用示例：
    #params_1064 = get_optical_parameters('1064nm')
    # params_810 = get_optical_parameters('810nm')
    # params_670 = get_optical_parameters('670nm')

    light_parameter = get_optical_parameters(wavelength)
    # light_parameter = np.array([
    #     [0, 0, 1, 1],# air
    #     [0.017, 18.45, 0.89, 1.37],# scalp
    #     [0.019, 14.6, 0.89, 1.37],# skull 
    #     [0.0144, 0.09, 0.89, 1.37],# csf
    #     [0.053, 5.9, 0.91, 1.37],#gm
    #     [0.105, 30, 0.88, 1.37],# wm
    #     #[0.033,9.35,0.89,1.37], # ECT
    #     # [0.105,30,0.88,1.37] # air cavities
    # ])
    def power2photon(d, p, t, length):
        h = 6.62607015e-34  # Planck constant
        c = 3e8  # Speed of light
        E = (h * c) / (length * 1e-9)  # Energy per photon
        N = (p * d * t) / E  # Number of photons
        return N
    N = power2photon(d, p, t, length)
    
    cfg = {
        'nphoton': timwin * N / (t * 60),
        'outputtype': 'energy',
        'vol': individual_atlas,
        'prop': light_parameter[[0, 1, 2, 3, 4,5], :].tolist(),
        'srcnum': 1,
        'srcpos': src_pos,#f3_subject_pos,#[100, 175, 215.0700],
        'srctype': 'pencil',
        'srcdir': src_dir,#[-0.5086, -0.1822, -0.8415],
        'issrcfrom0': 1,
        'tstart': 0,
        'tend': timwin,
        'tstep': timwin,
        'isspecular': 0,
        'isreflect': 1,
        'autopilot': 1,
        'gpuid': 1
    }
    if custom_src!=None:
        cfg.update(custom_src) 
    if stage_2_mode=="simple":
        cfg["nphoton"]=1e6

    
    print(cfg)
    # check source 
    show_src_vol(cfg["srcpos"],cfg["srcdir"],cfg["vol"])
    print(cfg["srcpos"],cfg["srcdir"],cfg["vol"].shape)
    
    # Run MCX simulation
    res = pmcx.mcxlab(cfg)
    flux = res['flux']
    
    view_result(res,cfg["vol"])
    view_result_mask(res,cfg["vol"],cfg["vol"]>2)
    
    original_img=nib.load(input_dict_v2["path"])
    original_affine = original_img.affine
    original_pixdim = original_img.header.get_zooms()
    flux_pixdim = np.ones_like(original_pixdim)  # Voxel dimensions for flux data
    zoom_factors = flux_pixdim / original_pixdim
    resized_flux = zoom(flux, zoom_factors, order=1)  # Use order=1 for linear interpolation
    resized_flux = resized_flux.astype(np.float32)
    resized_flux[resized_flux<1e-10]=1e-10
    resized_flux=np.log10(resized_flux)
    nii_image = nib.Nifti1Image(resized_flux, original_affine)
    if not os.path.exists(save_path): os.makedirs(save_path)
    np.save(os.path.join(save_path,"MCX_outputs"),input_dict_v2)
    nib.save(nii_image, os.path.join(save_path,"MCX_results_log.nii.gz"))
    nib.save(original_img,os.path.join(save_path,"MCX_input_vol.nii.gz"))
    print(f"MCX results saved to {save_path}")
    return {"res":res,"flux":flux,"resized_flux":resized_flux,"output_path":save_path}

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Run Monte Carlo simulation for photon transport in a head model.')
    parser.add_argument('--subject_path', required=True, help='Path to the the head model folder .')
    parser.add_argument('--save_path', required=True, help='Path to the output folder.')
    parser.add_argument('--region_name', required=True, help='Region name in the CSV file to be used as source position.')
    parser.add_argument('--p', type=float, required=True, help='Simulation parameter p.')
    parser.add_argument('--t', type=float, required=True, help='Simulation parameter t.')
    # 添加波长参数
    parser.add_argument('--wavelength', type=str, default='1064nm', 
                      choices=['670nm', '810nm', '850nm', '980nm', '1064nm'],
                      help='Wavelength for simulation (default: 1064nm)')

    args = parser.parse_args()
    subject_path = args.subject_path
    p = args.p
    t = args.t
    region_name = args.region_name
    save_path = args.save_path
    wavelength = args.wavelength  # 获取波长参数
    
    input_dict = {
        "path": os.path.join(subject_path,"final_tissues.nii.gz"),
        "subject_csv": os.path.join(subject_path,"eeg_positions","EEG10-10_UI_Jurak_2007.csv"),
        "seg_path": None,
        "region_name": region_name,
        "src_dir_mode": "default",
    }
    output_dict = run_fn(input_dict)
    vol = output_dict["vol"]
    src_dir = output_dict["src_dir"]
    src_pos = output_dict["src_pos"]
    output_dict["stage_2_mode"] = "full"
    output_dict["custom_src"] = None
    output_dict["save_path"] = save_path
    output_dict["p"] = p
    output_dict["t"] = t
    output_dict["wavelength"] = wavelength  # 添加波长参数到输出字典
    output_dict["path"]=input_dict["path"]
    output_dict_2 = run_fn_stage2(output_dict)
